GeoSAM2 Turns SAM2 Into a 3D Part Segmentation Tool

Analysis by the aitrendblend editorial team · Pillar: Vision transformers and attention · Source paper published August 2025

Why part segmentation in 3D has been stuck between slow and coarse

Getting a computer to say which triangles of a mesh belong to a chair leg versus a chair seat sounds like a small problem next to full scene understanding, yet it underpins robotic grasping, 3D asset generation, and interactive editing tools. The obvious route, train a supervised 3D network on labeled meshes, runs straight into a wall. Detailed part labels for 3D models take real annotator time to produce, so labeled datasets stay small compared to the flood of unlabeled 3D content being generated today.

That gap pushed researchers toward zero shot and weakly supervised methods that borrow strength from 2D vision foundation models instead. SAMPart3D distills multi view DINOv2 features into a 3D encoder and lets a user tune a continuous scale parameter to get finer or coarser parts. It works without text prompts, but the scale knob is a blunt instrument, nudging it up or down can produce splits that have nothing to do with actual part boundaries, and the method needs a fresh per shape optimization step that takes minutes per object. SAMesh lifts SAM masks from multiple rendered views and stitches them together with a community detection algorithm on the mesh graph, which avoids text prompts but offers no way to ask for a specific part and can take several minutes of iterative mesh optimization. PartField swings the other way, learning a continuous feature field in one fast feedforward pass, but users are again limited to picking a cluster count rather than pointing at the part they actually want.

Line those three up and a pattern appears. Methods are either fast and rigid or flexible and slow, and none of them let a user simply point at a part the way a human would say show me that leg. GeoSAM2 was built to close exactly that gap.

The core reframing, treat viewpoints like video frames

Instead of inventing a new 3D architecture from scratch, the authors lean on SAM2’s existing four stage design, an image encoder, a prompt and mask decoder, a memory attention module, and a memory encoder that compresses each predicted mask into a token for future frames to reference. That loop already handles promptable, temporally consistent segmentation across a video. GeoSAM2’s central move is to feed it something that is not a video at all.

Given a textureless 3D model, the pipeline renders normal maps and point maps from twelve canonical camera poses, evenly spaced in azimuth across three elevation angles at plus twenty five degrees, zero degrees, and minus twenty five degrees. Those twelve renders get sequenced counterclockwise around the object and handed to SAM2 as if they were consecutive video frames. A user annotates any single frame with a click or a box, and that frame becomes the sequence’s starting point. From there SAM2’s own tracking machinery carries the label around the object.

Because the object has no texture, a raw normal map alone leaves out useful spatial information, so each frame is paired with a point map built by back projecting the rendered depth into world coordinates.

Here \(D_i\) is the depth rendered from the mesh at view \(i\), \((u,v)\) is a pixel location, and \((R_i, \mathbf{t}_i)\) is the camera pose for that view. The resulting point map gives every pixel an actual 3D coordinate, which helps the model reason about correspondence across viewpoints that can differ far more sharply from each other than adjacent frames in a real video ever would.

Teaching a color model to read geometry without forgetting what it knows

SAM2 was pretrained on ordinary photographs, so its features are tuned for RGB appearance cues that simply do not exist in a normal map or a point map. Retraining the whole network on geometric inputs would risk erasing everything useful the pretraining already captured, and doing that kind of full fine tune at scale is expensive besides. The authors reach for Low-Rank Adaptation, freezing the base weight matrix of every relevant projection layer and adding a small trainable correction alongside it.

For an input feature vector \(f\), the layer output becomes \(W_0 f + A(Bf)\), with \(W_0\) held fixed and only the low rank matrices \(A\) and \(B\) trained. Two separate LoRA branches are attached, one tuned for normal map inputs and one for point map inputs, inserted into the query, key, and value projections of every transformer block. Only rank four adapters are needed per branch, a tiny fraction of the parameters in SAM2’s full backbone, yet the paper’s own ablation shows this single change is enough to bridge most of the gap between photometric and geometric inputs.

Fusing two geometric signals without shocking the network

Normal maps end up closer in statistics to the RGB textures SAM2 already understands, so they carry most of the early signal. Point maps carry the spatial correspondence information the model actually needs, but injecting them carelessly risks knocking the pretrained feature statistics out of their comfortable range. The fusion scheme handles that by starting from nothing and letting the network earn its way in.

The normal feature \(\mathbf{G}_i\) and the point map feature \(\mathbf{P}_i\) are concatenated along the channel axis and passed through a three by three convolution whose weights start at exactly zero. At the very first training step that convolution outputs nothing, so the fused feature \(\hat{\mathbf{G}}_i\) is identical to the plain normal feature the network already knows how to process. As training proceeds, gradients slowly open the door for the point map branch to contribute, rather than forcing it in on day one. This same trick, zero initialized convolutions guarding a residual path, shows up across a lot of recent adapter and control style architectures, and it earns its keep here for the same reason it does elsewhere, stability during the first few thousand steps of training.

A memory bank redesigned for viewpoints instead of time

SAM2’s original memory bank is a fixed size, first in first out queue, a sensible design when adjacent video frames are nearly redundant and old frames stop being useful once the scene has moved on. That assumption breaks completely in a twelve view sequence. Each of the twelve renders carries a genuinely different, non redundant slice of the object’s geometry, so quietly discarding the third view once the model reaches the ninth throws away information no later frame can recover. GeoSAM2 simply retains every view’s embedding for the full sequence rather than letting a queue evict them.

A second, smaller fix turned out to matter more than expected. The very first frame in any sequence starts with an empty memory bank and only a sparse click or box to work from, and the resulting mask quality on that frame alone is noticeably weaker. Once the memory bank has even one populated entry, quality jumps. Rather than solving that with a more complicated architecture, the authors duplicate the first frame before the sequence begins, giving the model an immediate, meaningful memory prior instead of a blank one. It is a small trick with a clear before and after, shown in the paper’s own frame repetition ablation, where masks with the duplicated frame track edges more tightly and stay more stable across views.

From twelve masks to one labeled mesh

Once every view has a predicted 2D mask, the pipeline back projects each one onto the 3D surface using the known camera poses. For mesh inputs specifically, five points are sampled per face and projected into every view. A depth comparison check, comparing the mesh’s own rendered depth against each mask’s depth map, discards points whose depth mismatch exceeds a small threshold, filtering out points that are actually occluded from that particular view rather than truly visible and mislabeled. Each visible point then takes the majority label across the views that can actually see it.

A lightweight cleanup pass follows. Small disconnected components below an area threshold get removed, and remaining unlabeled faces are filled in using k nearest neighbor voting against already labeled faces, guaranteeing every face on the mesh ends up with a label rather than leaving gaps.

How it stacks up against the field

The authors compare against four recent open world 3D part segmentation methods, evaluated on their public implementations with the same experimental setup used in the PartField paper. Find3D is text prompted rather than class agnostic, SAMPart3D needs per shape optimization on top of its 3D pretraining, SAMesh depends on mesh connectivity for its community detection step, and PartField predicts a triplane feature field with clustering based granularity control.

PartObjaverse-Tiny, 200 human labeled meshes

The text prompted Find3D trailing every class agnostic method by a wide margin is itself a small finding, it suggests that tying part discovery to open vocabulary text is still a harder problem than simply discovering parts and letting a user point at the one they want. SAMesh’s acceptable average score hides a habit the authors flag directly, it tends to over segment a mesh into pieces finer than what a human would call a part, with no way for a user to dial that back.

PartNetE, 1,906 point clouds across 45 categories

The move from PartObjaverse-Tiny to PartNetE is where the comparison gets genuinely informative. SAMesh and PartField both drop twenty to thirty points of mIoU on this second benchmark, which the authors trace directly to a dependency on mesh connectivity, since PartNetE is distributed as point clouds rather than clean meshes. Find3D and SAMPart3D barely move between the two benchmarks, because their post processing already leans on k nearest neighbor voting rather than mesh topology. GeoSAM2 also loses some ground without mesh connectivity to lean on, dropping from 84.06 to 74.42, but it keeps a wide lead over every baseline in both settings, and it does so while running in roughly thirty seconds, far faster than SAMPart3D or SAMesh and only modestly slower than the two feedforward methods.

What the ablation study isolates

The team peeled the method back one layer at a time, starting from vanilla SAM2 fed multi view normal maps with no adaptation at all, then adding LoRA tuning for normal maps only, then adding the point map branch through a plain one by one convolution with no residual fusion, and finally the complete method.

Vanilla SAM2 already manages a respectable 62.59 on PartObjaverse-Tiny purely on the strength of its pretrained tracking ability, but the paper notes its failures come from an inability to precisely track without any geometric grounding at all. Adding LoRA tuned normal map processing closes a large chunk of that gap by itself, roughly thirteen points, confirming that most of the domain shift between photographic and geometric inputs can be handled by a lightweight adapter rather than a full retrain. The point map branch adds spatial disambiguation on top of that, and the final residual fusion step is what lets the model hold onto fine detail that a plain convolution without the residual path tends to blur away.

Where the idea travels beyond a single mesh

Because the whole pipeline is prompt driven rather than tied to a fixed label set, the authors show it working on several downstream tasks without any retraining. Paired with a 3D part completion model such as HoloPart, GeoSAM2 can perform 3D part amodal segmentation, recovering the hidden geometry of a part that continues behind another surface. Applied to meshes generated by TripoSG, a diffusion based 3D generator, the method segments hierarchical part structures, hands, legs, and head at varying levels of detail, even though generated meshes rarely have the clean, artist made connectivity that real assets do. For interactive 3D editing, a single click can isolate something as specific as a human leg or a patch of a beard, with the boundary propagating cleanly to the full mesh, which the authors position as a real alternative to fiddling with a scale slider or writing a text prompt.

Honest limitations

Complete PyTorch implementation

The code below is an independent, runnable reimplementation of GeoSAM2’s core building blocks, the LoRA adapted linear layer, the point map projection, the zero initialized residual fusion module, and a simplified multi view mask aggregation step, built directly from the equations and architecture description in the paper. It favors clarity over production performance and ends with a smoke test on random tensors.

# geosam2_core.py
# Independent PyTorch reimplementation of the core GeoSAM2 components
# Based on Deng, Yang, Sun, Liu, Liu, Liang, and Cao, arXiv:2508.14036

import torch
import torch.nn as nn
import torch.nn.functional as F


class LoRALinear(nn.Module):
    """A frozen linear layer with a trainable low rank correction.

    Matches the paper's LoRA formulation, W = W0 + AB, where W0 stays
    frozen and only the small matrices A and B receive gradients.
    """

    def __init__(self, in_features, out_features, rank=4):
        super().__init__()
        self.base = nn.Linear(in_features, out_features, bias=False)
        self.base.weight.requires_grad = False
        self.lora_a = nn.Parameter(torch.randn(in_features, rank) * 0.01)
        self.lora_b = nn.Parameter(torch.zeros(rank, out_features))

    def forward(self, x):
        base_out = self.base(x)
        lora_out = (x @ self.lora_a) @ self.lora_b
        return base_out + lora_out


class DualLoRAAttentionProjection(nn.Module):
    """Two independent LoRA branches sharing one frozen base projection.

    One branch specializes in normal map features, the other in point
    map features, matching the paper's per modality adapter design.
    """

    def __init__(self, dim, rank=4):
        super().__init__()
        self.shared_base = nn.Linear(dim, dim, bias=False)
        self.shared_base.weight.requires_grad = False
        self.normal_lora = LoRALinear(dim, dim, rank=rank)
        self.point_lora = LoRALinear(dim, dim, rank=rank)

    def forward(self, x, modality):
        if modality == "normal":
            return self.normal_lora(x)
        elif modality == "point":
            return self.point_lora(x)
        raise ValueError("modality must be normal or point")


def compute_point_map(depth, k_inv, r_inv, t):
    """Back projects a rendered depth map into world space point coordinates.

    Mirrors the paper's point map equation, giving each pixel a 3D
    coordinate derived from the camera intrinsics and extrinsics.
    """
    batch, height, width = depth.shape
    ys, xs = torch.meshgrid(
        torch.arange(height, device=depth.device),
        torch.arange(width, device=depth.device),
        indexing="ij",
    )
    ones = torch.ones_like(xs)
    pixels = torch.stack([xs, ys, ones], dim=0).float()
    pixels = pixels.view(3, -1)

    camera_rays = k_inv @ pixels
    camera_rays = camera_rays.view(1, 3, height, width).expand(batch, -1, -1, -1)

    scaled = depth.unsqueeze(1) * camera_rays
    world = torch.einsum("ij,bjhw->bihw", r_inv, scaled) - (r_inv @ t).view(1, 3, 1, 1)
    return world


class ZeroInitResidualFusion(nn.Module):
    """Fuses normal and point map features via a zero initialized convolution.

    At initialization the fused output equals the plain normal feature,
    matching the paper's stability trick for introducing a new modality
    without shocking the pretrained backbone.
    """

    def __init__(self, channels):
        super().__init__()
        self.fuse = nn.Conv2d(channels * 2, channels, kernel_size=3, padding=1)
        nn.init.zeros_(self.fuse.weight)
        nn.init.zeros_(self.fuse.bias)

    def forward(self, normal_feat, point_feat):
        fused_input = torch.cat([normal_feat, point_feat], dim=1)
        residual = self.fuse(fused_input)
        return normal_feat + residual


class MultiViewMemoryBank(nn.Module):
    """Retains every view's embedding rather than a fixed size FIFO queue.

    Also supports duplicating the first frame before the real sequence
    begins, matching the paper's bootstrapping strategy.
    """

    def __init__(self):
        super().__init__()
        self.tokens = []

    def reset(self):
        self.tokens = []

    def bootstrap(self, first_frame_token):
        self.tokens = [first_frame_token]

    def add(self, token):
        self.tokens.append(token)

    def stacked(self):
        return torch.stack(self.tokens, dim=1) if self.tokens else None


def aggregate_multiview_labels(view_labels, visibility_mask):
    """Combines per view point labels into one label per 3D point.

    view_labels has shape views by points, visibility_mask is a
    matching boolean tensor marking which views can actually see each
    point after the depth consistency check described in the paper.
    """
    num_views, num_points = view_labels.shape
    num_classes = int(view_labels.max().item()) + 1

    votes = torch.zeros(num_points, num_classes)
    for v in range(num_views):
        visible_points = visibility_mask[v].nonzero(as_tuple=True)[0]
        labels_here = view_labels[v, visible_points]
        votes[visible_points, labels_here] += 1

    final_labels = votes.argmax(dim=1)
    return final_labels


if __name__ == "__main__":
    # Smoke test on random dummy data, confirms each building block runs
    # before wiring in a real SAM2 backbone and a real renderer.
    torch.manual_seed(0)

    # LoRA branch sanity check
    dual = DualLoRAAttentionProjection(dim=32, rank=4)
    tokens = torch.randn(2, 10, 32)
    normal_out = dual(tokens, "normal")
    point_out = dual(tokens, "point")
    print("normal branch output shape", normal_out.shape)
    print("point branch output shape", point_out.shape)

    # Point map construction sanity check
    depth = torch.rand(1, 16, 16) + 0.5
    k_inv = torch.eye(3)
    r_inv = torch.eye(3)
    t = torch.zeros(3, 1)
    point_map = compute_point_map(depth, k_inv, r_inv, t)
    print("point map shape", point_map.shape)

    # Residual fusion sanity check, confirms zero init means identity at start
    fusion = ZeroInitResidualFusion(channels=8)
    normal_feat = torch.randn(1, 8, 16, 16)
    point_feat = torch.randn(1, 8, 16, 16)
    fused = fusion(normal_feat, point_feat)
    identical_at_init = torch.allclose(fused, normal_feat)
    print("fusion output equals normal features at init", identical_at_init)

    # Multi view memory bank sanity check
    bank = MultiViewMemoryBank()
    bank.bootstrap(torch.randn(4))
    for _ in range(11):
        bank.add(torch.randn(4))
    print("memory bank size after full sequence", len(bank.tokens))

    # Multi view label aggregation sanity check
    view_labels = torch.randint(0, 3, (12, 100))
    visibility = torch.rand(12, 100) > 0.4
    visibility[0] = True
    final_labels = aggregate_multiview_labels(view_labels, visibility)
    print("aggregated label shape", final_labels.shape)

Running this file end to end should print shapes and a confirmation that the fusion block behaves as an identity function on the normal features before any training happens, which is the entire point of the zero initialization trick.

Conclusion

GeoSAM2’s biggest contribution is not a bigger model or a new loss function, it is a genuinely clever piece of reuse. SAM2 already knew how to take a sparse prompt on one frame and carry it consistently through a sequence, and the authors recognized that a spinning object rendered from twelve angles is, structurally, just another sequence, provided you are willing to rebuild the parts of SAM2 that assumed time rather than viewpoint was the thing connecting one frame to the next.

That is also the conceptual shift worth remembering past this specific paper. A foundation model’s architecture often encodes assumptions that are more specific than they first appear, SAM2’s memory bank was not just any memory mechanism, it was a memory mechanism tuned for temporal redundancy. Recognizing which assumption needs surgery, and which parts of the pretrained network can stay frozen, is a more transferable skill than any single number in the results tables.

The idea should travel reasonably well beyond meshes and figurines. Anything that can be rendered from multiple canonical viewpoints, architectural scans, manufactured parts, medical phantoms used for non clinical research, could plausibly be run through the same multi view video reframing, provided a normal map and a point map can be produced for it. The LoRA and zero init fusion recipe is not tied to any particular kind of 3D geometry.

The limitations are worth taking as seriously as the wins. Occlusion remains a real weakness for any method built on stitching outside views together, mesh connectivity still provides a meaningful boost that point cloud only inputs cannot supply, and the training set behind this specific model is a modest 4,700 self annotated objects rather than an internet scale corpus. None of that erases what the paper demonstrates, that a two year old video segmentation model still has capacity nobody had extracted yet, if you are willing to rethink what a frame is.

As more 3D generation pipelines ship models with messy, non artist grade topology, tools that segment by pointing rather than by tuning a scale parameter are likely to matter more, not less. GeoSAM2 is a solid, thoroughly ablated argument for building those tools on top of a video foundation model rather than starting a 3D architecture from zero.

Frequently asked questions

Related reading